package com.baidu.encrypt.common.utils;

import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.paddings.PKCS7Padding;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.security.Security;
import java.util.Base64;

@Component
public class SM4Util {

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    @Value("${sm4.key}")
    private String keyConfig;

    @Value("${sm4.iv}")
    private String ivConfig;

    private byte[] keyBytes;
    private byte[] ivBytes;

    @PostConstruct
    public void init() {
        this.keyBytes = keyConfig.getBytes(StandardCharsets.UTF_8);
        this.ivBytes = ivConfig.getBytes(StandardCharsets.UTF_8);
    }

    /** 字符串加密 */
    public String encrypt(String plainText) throws Exception {
        byte[] encrypted = process(true, plainText.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(encrypted);
    }

    /** 字符串解密 */
    public String decrypt(String cipherText) throws Exception {
        byte[] decrypted = process(false, Base64.getDecoder().decode(cipherText));
        return new String(decrypted, StandardCharsets.UTF_8);
    }

    /** 流式加密 */
    public void encryptStream(InputStream in, OutputStream out) throws Exception {
        processStream(true, in, out);
    }

    /** 流式解密 */
    public void decryptStream(InputStream in, OutputStream out) throws Exception {
        processStream(false, in, out);
    }

    /** 核心字节处理方法 */
    private byte[] process(boolean forEncryption, byte[] input) throws Exception {
        PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(
                new CBCBlockCipher(new SM4Engine()), new PKCS7Padding());
        cipher.init(forEncryption, new ParametersWithIV(new KeyParameter(keyBytes), ivBytes));

        byte[] output = new byte[cipher.getOutputSize(input.length)];
        int len = cipher.processBytes(input, 0, input.length, output, 0);
        len += cipher.doFinal(output, len);

        byte[] result = new byte[len];
        System.arraycopy(output, 0, result, 0, len);
        return result;
    }

    /** 核心流处理方法 */
    private void processStream(boolean forEncryption, InputStream in, OutputStream out) throws Exception {
        PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(
                new CBCBlockCipher(new SM4Engine()), new PKCS7Padding());
        cipher.init(forEncryption, new ParametersWithIV(new KeyParameter(keyBytes), ivBytes));

        byte[] inBuffer = new byte[1024];
        byte[] outBuffer = new byte[cipher.getOutputSize(inBuffer.length)];
        int inBytes;
        while ((inBytes = in.read(inBuffer)) != -1) {
            int outBytes = cipher.processBytes(inBuffer, 0, inBytes, outBuffer, 0);
            if (outBytes > 0) {
                out.write(outBuffer, 0, outBytes);
            }
        }
        int outBytes = cipher.doFinal(outBuffer, 0);
        if (outBytes > 0) {
            out.write(outBuffer, 0, outBytes);
        }
        out.flush();
    }
}
